# -*- coding: UTF-8 -*-
'''
@Project ：push_rk
@File ：myhandle.py
@IDE ：PyCharm
@Author ：苦瓜
@Date ：2025/10/9 10:33
@Note: Something beautiful is about to happen !
'''
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn

# 1.	加载CSV文件中的数据，并反转数据以便按照时间顺序排列
data = np.loadtxt("data-02-stock_daily.csv", delimiter=",")
# 2.	对数据进行归一化处理
from sklearn.preprocessing import MinMaxScaler
data = MinMaxScaler().fit_transform(data)
# 3.	设置时间步长
c = 7
# 4.	初始化输入和输出数据列表
x = []
y = []

# 5.	遍历数据，构建输入和输出
for i in range(len(data) - c):
    x.append(data[i:i+c])
    y.append(data[i+c, -1:])

# 6.	将输入和输出数据转换为Tensor
X = torch.tensor(x, dtype=torch.float32)
Y = torch.tensor(y, dtype=torch.float32)
# 7.	划分数据集为训练集和测试集
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, Y, shuffle=False, random_state=42)
# 8.	打印训练数据的形状
print(f"""
    ===  切分后的数据形状  ===
    X_train shape : {X_train.shape}
    X_test shape : {X_test.shape}
    y_train shape : {y_train.shape}
    y_test shape : {y_test.shape}
""")
# 9.	设置输入数据的特征维度
input_dim = X_train.shape[2]
# 10.	定义RNN模型
# RNN 为关键字 使用下划线规范化
class RNN_(nn.Module):

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.rnn1 = nn.RNN(input_size=input_dim, hidden_size=10, batch_first=True)
        self.rnn2 = nn.RNN(input_size=10, hidden_size=5, batch_first=True)
        self.fc = nn.Linear(in_features=5, out_features=1)

    def forward(self, inputs):
        x_, *_ = self.rnn1(inputs)
        x_, *_ = self.rnn2(x_)
        x_ = self.fc(x_[:, -1, :])
        return x_

if __name__ == '__main__':
    # 11.	实例化模型
    model_rnn = RNN_()
    # 12.	定义损失函数
    cost = nn.MSELoss()
    # 13.	定义优化器
    optimizer = torch.optim.Adam(model_rnn.parameters(), lr=0.001)
    # 14.	初始化损失列表
    losses_ = []
    # 15.	训练模型
    model_rnn.train()
    for epoch in range(1000):
        optimizer.zero_grad()
        y_pred_train = model_rnn(X_train)
        loss_ = cost(y_pred_train, y_train)
        losses_.append(loss_.item())
        loss_.backward()
        optimizer.step()
        if epoch % 10 == 0:
            print(f"<Model RNN_ train loss : ({loss_})>")

    model_rnn.eval()
    # 16.	模型预测测试集
    with torch.no_grad():
        y_pred_test = model_rnn(X_test)

    # 17.	打印预测结果
    print(y_pred_test)
    # 18.	绘制预测结果和真实结果
    plt.plot(y_pred_test, label='y pred test', c='g')
    plt.plot(y_test, label='y test', c='r')
    plt.legend()
    plt.show()

    # 损失图
    plt.plot(losses_, label='loss')
    plt.legend()
    plt.show()

